# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Callback to write H5MD trajectory file
"""

import signal
import os
import sys
from typing import List
import numpy as np
import torch

from .callback_manager import Callback, RunContext

def check_non_negative_int(value):
    """Check if value is a non-negative integer"""
    if not isinstance(value, int) or value < 0:
        raise ValueError(f"Value must be a non-negative integer, but got: {value}")
    return value

def check_bool(value):
    """Check if value is a boolean"""
    if not isinstance(value, bool):
        raise ValueError(f"Value must be a boolean, but got: {value}")
    return value

from ..system import Molecule
from ..optimizer import Updater
from ..data.export import H5MD
from ..system.modelling.pdb_generator import gen_pdb


class WriteH5MD(Callback):
    r"""Callback to write HDF5 molecular data (H5MD) file

    Args:
        system (Molecule):      Simulation system (system module of MindSPONGE)
        filename (str):         Name of output H5MD file
        directory (str):        Directory of the output file. Default: ``None``
        mode (str):            I/O mode for H5MD
                                - ``'w'``:            Create file, truncate if exists
                                - ``'w-'`` or ``'x'``: Create file, fail if exists
                                - ``'a'``:            Read/write if exists, create otherwise
        write_velocity (bool):  Whether to write the velocity of the system to the H5MD file.
                                Default: ``False``
        write_force (bool):     Whether to write the force of the system to the H5MD file.
                                Default: ``False``
        write_image (bool):     Whether to write the image of the position of system to the H5MD file.
                                Default: ``False``
        length_unit (str):      Length unit for coordinates. Default: ``None``
        energy_unit (str):      Energy unit. Default: ``None``
        dtype (str):            Data type for H5MD. Default: ``'float32'``
        compression (str):      Compression strategy for HDF5. Default: ``'gzip'``
        compression_opts (int): Compression settings for HDF5. Default: ``4``
        auto_close (bool):      Whether to automatically close the writing of H5MD files at the end of
                                the simulation process. Default: ``True``
        save_last_pdb (str):    Whether to store the last coordinates in a pdb format file. If choose to store the pdb,
                                the value should be string format pdb file name. Default: ``None``
        batch_index (int):      The batch index of system for H5MD file to store.
                                Default: ``None``

    Supported Platforms:
        ``CPU`` ``GPU``
    """

    def __init__(self,
                 system: Molecule,
                 filename: str,
                 save_freq: int = 1,
                 directory: str = None,
                 mode: str = 'w',
                 write_velocity: bool = False,
                 write_force: bool = False,
                 write_image: bool = False,
                 write_metrics: bool = None,
                 write_energies: bool = False,
                 write_bias: bool = False,
                 length_unit: str = None,
                 energy_unit: str = None,
                 dtype: str = 'float32',
                 compression: str = 'gzip',
                 compression_opts: int = 4,
                 auto_close: bool = True,
                 save_last_pdb: str = None,
                 batch_index: int = None
                 ):

        if mode not in ['w', 'w-', 'x', 'a']:
            raise ValueError(f'[WriteH5MD] The supported types of mode are "w", "w-", "x" and "a", '
                             f'but got: {mode}')

        self.system = system
        self.units = system.units
        self.batch_index = batch_index
        if batch_index is not None:
            self.num_walker = 1
        else:
            self.num_walker = system.num_walker
        if mode == 'w' or mode == 'w-':
            if os.path.exists(filename):
                os.remove(filename)
        self.h5md = H5MD(self.system, filename, directory, mode=mode,
                         length_unit=length_unit, energy_unit=energy_unit,
                         compression=compression, compression_opts=compression_opts, 
                         num_walker=self.num_walker)
        self.filename = filename

        self.convert_to_angstram = self.system.units.convert_length_to('A')
        self.use_pbc = system.pbc_box is not None
        self.const_volume = True
        self.dtype = dtype
        self.save_freq = check_non_negative_int(save_freq)
        self.auto_close = check_bool(auto_close)

        self.write_image = check_bool(write_image)
        self.write_velocity = check_bool(write_velocity)
        self.write_force = check_bool(write_force)

        if save_last_pdb is None:
            self.last_pdb_name = None
            self.save_pdb = False
        else:
            self.save_pdb = True
            self.last_pdb_name = save_last_pdb

            if os.path.exists(self.last_pdb_name):
                os.remove(self.last_pdb_name)

        if mode == 'a':
            self.init_h5md = False
        else:
            self.init_h5md = True
            if self.use_pbc:
                self.h5md.set_box(False,True,True)
                if self.write_image:
                    self.h5md.set_image()

            if self.write_velocity:
                self.h5md.set_velocity()

            if self.write_force:
                self.h5md.set_force()

            self.h5md.add_observables('potential_energy', (), self.dtype, self.units.energy_unit_name)
            self.h5md.add_observables('total_energy', (), self.dtype, self.units.energy_unit_name)

        self.use_updater = None

        self.write_metrics = write_metrics
        self.write_energies = write_energies
        self.write_bias = write_bias
        self.metad = False

        self.count = 0
        self.count_records = 0

        # Detect process kill signal and save the data
        signal.signal(signal.SIGINT, self.__stop__)
        signal.signal(signal.SIGTERM, self.__stop__)

    def __enter__(self):
        """Return the enter target"""
        return self

    def __exit__(self, *err):
        """Release resources here if have any"""

    def __stop__(self, signal_, frame_):
        """
        Save data when process killed
        """
        # pylint: disable=unused-argument
        print(f'\n\033[33mProgram process terminated. {self.count_records} steps saved in H5MD file.\033[0m\n')
        self.close()
        sys.exit(0)

    def close(self):
        """Close the H5MD file"""
        self.h5md.close()
        return self

    def begin(self, run_context: RunContext):
        """
        Called once before the network executing

        Args:
            run_context (RunContext): Include some information of the model
        """
        
        print(f'[TorchSPONGE] [WriteH5MD] Writing H5MD file: {self.filename}', )
        # Check if h5md file is open
        self.init_h5md = self.h5md.check_open()

        if not self.init_h5md:
            self.h5md = self.h5md.reload()
            return
        
        if isinstance(run_context.optimizer, Updater):
            self.use_updater = True
            self.h5md.add_observables('kinetic_energy', (), self.dtype, self.units.energy_unit_name)
            self.h5md.add_observables('temperature', (), self.dtype, 'K')
            if self.use_pbc:
                self.h5md.add_observables('pressure', (), self.dtype, 'bar')
                self.h5md.add_observables('volume', (), self.dtype, 'nm^3')

        if self.write_metrics:
            for k,v in run_context.metrics.items():
                self.h5md.add_observables(k, (), self.dtype, None)
        if self.write_energies:
            self._init_energies(run_context.num_energies, run_context.energy_names)
        if self.write_bias:
            self._init_biases(run_context.bias_names, run_context.bias_function)

    def _init_energies(self, energy_names: List[str]):
        """Initialize energy observables"""
        self.write_energies = True
        num_energies = len(energy_names)
        energies = self.h5md.add_observables('energies', (num_energies,), self.dtype, self.units.energy_unit_name)
        self.h5md.create_dataset(energies, 'labels', shape=(num_energies,), data=energy_names)

    def _init_biases(self, bias_names: List[str], bias_function:List):
        """Initialize bias observables"""
        self.write_bias = True
        num_biases = len(bias_names)
        biases = self.h5md.add_observables('biases', (num_biases,), self.dtype, self.units.energy_unit_name)
        self.h5md.add_observables('bias_potential', (), self.dtype, self.units.energy_unit_name)
        self.h5md.create_dataset(biases, 'labels', shape=(num_biases,), data=bias_names)
        for i in range(num_biases):
            if 'metadynamics' in bias_names[i].lower():
                self.h5md.create_metadynamics(bias_function[i].hills.detach().cpu().numpy(), bias_function[i].grids.detach().cpu().numpy())
                self.metad = True
                self.metad_idx = i
                
    def epoch_begin(self, run_context: RunContext):
        """
        Called before each epoch beginning

        Args:
            run_context (RunContext): Include some information of the model
        """

    def epoch_end(self, run_context: RunContext):
        """
        Called after each epoch finished

        Args:
            run_context (RunContext): Include some information of the model
        """

    def step_begin(self, run_context: RunContext):
        """
        Called before each step beginning

        Args:
            run_context (RunContext): Include some information of the model
        """

    def step_end(self, run_context: RunContext):
        """
        Called after each step finished

        Args:
            run_context (RunContext): Include some information of the model
        """
        if self.count % self.save_freq == 0:
            self.call_end(run_context)
            self.count_records += 1

        self.count += 1

    def call_begin(self, run_context: RunContext):
        """
        Called at the beginning

        Args:
            run_context (RunContext): Include some information of the model
        """

    def call_end(self, run_context: RunContext):
        """
        Called at the end

        Args:
            run_context (RunContext): Include some information of the model
        """
        if self.batch_index is not None:
            coordinate = run_context.coordinate[self.batch_index]
            pbc_box = run_context.pbc_box[self.batch_index]
            if self.write_velocity:
                velocity = run_context.velocity[self.batch_index]
            if self.write_force:
                force = run_context.force[self.batch_index]
            if self.use_pbc and self.write_image:
                image = run_context.image[self.batch_index]
            potential = run_context.potential[self.batch_index]
            if self.use_updater:
                kinetics = run_context.kinetics[self.batch_index]
                temperature = run_context.temperature[self.batch_index]
                if self.use_pbc:
                    pressure = run_context.pressure[self.batch_index]
                    volume = run_context.volume[self.batch_index]
            if self.write_metrics:
                metrics = {}
                for k,v in run_context.metrics.items():
                    metrics[k] = metrics[v].eval()[self.batch_index]
            if self.write_energies:
                energies = run_context.energies[self.batch_index]
            if self.write_bias:
                biases = run_context.biases[self.batch_index]
                bias = run_context.bias[self.batch_index]

        else:
            coordinate = run_context.coordinate
            pbc_box = run_context.pbc_box
            if self.write_velocity:
                velocity = run_context.velocity
            if self.write_force:
                force = run_context.force
            if self.use_pbc and self.write_image:
                image = run_context.image
            potential = run_context.potential
            if self.use_updater:
                kinetics = run_context.kinetics
                temperature = run_context.temperature
                if self.use_pbc:
                    pressure = run_context.pressure
                    volume = run_context.volume
                tot_energy = potential + kinetics
            else:
                tot_energy = potential
            if self.write_metrics:
                metrics = {}
                for k,v in run_context.metrics.items():
                    metrics[k] = v.eval()
            if self.write_energies:
                energies = run_context.energies
            if self.write_bias:
                biases = run_context.biases
                bias = run_context.bias

        step = run_context.cur_step
        time = run_context.cur_time

        self.h5md.write_position(coordinate, step, time)
        if self.write_velocity:
            self.h5md.write_velocity(velocity, step, time)
        if self.write_force:
            self.h5md.write_force(force, step, time)
        if self.use_pbc:
            self.h5md.write_box(pbc_box, step, time)
            if self.write_image:
                self.h5md.write_image(image, step, time)

        self.h5md.write_observables('potential_energy', potential, step, time)
        self.h5md.write_observables('total_energy', tot_energy, step, time)

        if self.use_updater:
            self.h5md.write_observables('kinetic_energy', kinetics, step, time)
            self.h5md.write_observables('temperature', temperature, step, time)
            if self.use_pbc:
                self.h5md.write_observables('pressure', pressure, step, time)
                self.h5md.write_observables('volume', volume, step, time)

        if self.write_metrics:
            for k,v in metrics.items():
                self.h5md.write_observables(k, v, step, time)
        if self.write_energies:
            self.h5md.write_observables('energies', energies, step, time)
        if self.write_bias:
            self.h5md.write_observables('bias_potential', bias, step, time)
            self.h5md.write_observables('biases', biases, step, time)
            if self.metad:
                self.h5md.update_metadynamics(run_context.bias_function[self.metad_idx].hills.detach().cpu().numpy())
                    
        self.h5md.update_restart(coordinate,pbc_box)
        self.h5md.hdf5_file.flush()

    def end(self, run_context: RunContext):
        """
        Called once after network training

        Args:
            run_context (RunContext): Include some information of the model
        """
        if self.auto_close:
            self.close()
        if self.save_pdb:
            self.save_to_pdb()

    def save_to_pdb(self):
        """Save the system information into a pdb file"""
        last_resname = self.system.residue_name
        for i, name in enumerate(last_resname):
            last_resname[i] = name[-3:]

        if torch.isnan(self.system.coordinate.sum()):
            return None

        # Clear the pdb path
        if os.path.exists(self.last_pdb_name):
            os.remove(self.last_pdb_name)

        coordinate = self.system.coordinate.detach().cpu().numpy()[0] * self.convert_to_angstram
        atom_resid = self.system.atom_resid.detach().cpu().numpy()

        gen_pdb(coordinate,
                self.system.atom_name[0],
                np.take(last_resname, self.system.atom_resid),
                atom_resid + 1,
                pdb_name=self.last_pdb_name)